import torch
import torch.nn as nn


import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt



torch.manual_seed(20)

# 学习率
lr = 0.001

x = torch.rand(20, 1)*10
y = 2 * x + (5 + torch.randn(20, 1))

w = torch.randn((1), requires_grad=True)
b = torch.randn((1), requires_grad=True)

for iteration in range(100000):
    y_pred = w * x + b

    # loss = (0.5 * (y - y_pred) ** 2).mean()
    loss = nn.MSELoss()(y_pred, y)

    loss.backward()

    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad

    # 清零梯度
    w.grad.zero_()
    b.grad.zero_()

    if iteration % 30 == 0:
        plt.cla()

        plt.scatter(x.data.numpy(), y.data.numpy())
        plt.plot(x.data.numpy(), y_pred.data.numpy(), 'r', lw=5)
        plt.text(2, 20, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'red'})
        plt.xlim(1.5, 10)
        plt.ylim(8, 28)

        plt.title('Iteration: {}\n w:{} b:{} '.format(iteration, w.data.numpy(), b.data.numpy()))

        plt.pause(0.01)

        if loss.data.numpy() < 0.5:
            break

plt.show()
